import numpy as np
from numba import cuda
from numba.cuda.cudadrv.devicearray import DeviceNDArray


class Tensor:
    def __init__(self, shape, dtype=float, strides=None, data=None, g_data=None):
        self.shape = tuple(shape)
        self.dtype = dtype
        self.strides = tuple(strides)
        self.data: np.ndarray = data
        self.g_data: DeviceNDArray = g_data

    def to_device(self):
        assert self.data is not None

        if self.g_data is None:
            self.g_data = cuda.device_array(self.shape, dtype=self.dtype, strides=self.strides, order='C', stream=0)
        self.g_data.copy_to_device(self.data)

        return self.g_data

    def to_host(self):
        assert self.g_data is not None

        if self.data is None:
            self.data = np.ndarray(self.shape, dtype=self.dtype, strides=self.strides, order='C')
        self.g_data.copy_to_host(self.data)
        return self.data

    def __str__(self):
        return self.data.__str__()


def from_ndarray(arr, host=True, device=False):
    shape, dtype, strides = arr.shape, arr.dtype, arr.strides

    g_data = None
    if device:
        g_data = cuda.device_array(shape, dtype=dtype, strides=strides, order='C', stream=0)
        g_data.copy_to_device(arr)

    data = None
    if host:
        data = arr
    else:
        del arr

    return Tensor(shape, dtype=dtype, strides=strides, data=data, g_data=g_data)


def array(p_object, dtype=None, host=True, device=False):
    arr = np.array(p_object, dtype=dtype)
    return from_ndarray(arr, host=host, device=device)


def arange(start=None, stop=None, step=None, dtype=None, host=True, device=False):
    arr = np.arange(start, stop, step, dtype)
    return from_ndarray(arr, host=host, device=device)


def empty(shape, dtype=None, host=True, device=False):
    arr = np.empty(shape, dtype)
    return from_ndarray(arr, host=host, device=device)


def empty_like(a, dtype=None, host=True, device=False):
    arr = np.empty_like(a, dtype)
    return from_ndarray(arr, host=host, device=device)


def ones(shape, dtype=None, host=True, device=False):
    arr = np.ones(shape, dtype)
    return from_ndarray(arr, host=host, device=device)


def ones_like(a, dtype=None, host=True, device=False):
    arr = np.ones_like(a, dtype)
    return from_ndarray(arr, host=host, device=device)


def zeros(shape, dtype=None, host=True, device=False):
    arr = np.zeros(shape, dtype)
    return from_ndarray(arr, host=host, device=device)


def zeros_like(a, dtype=None, host=True, device=False):
    arr = np.zeros_like(a, dtype)
    return from_ndarray(arr, host=host, device=device)


def eye(N, M=None, k=0, dtype=float, host=True, device=False):
    arr = np.eye(N, M, k, dtype)
    return from_ndarray(arr, host=host, device=device)
